"""
DeepSeek-R1 模型辅助类
此模块提供了与DeepSeek-R1模型交互的功能
"""

import os
import json
import requests

class DeepSeekR1:
    def __init__(self, api_key=None):
        """
        初始化DeepSeek-R1助手
        
        参数:
            api_key (str, 可选): DeepSeek API密钥，使用API时需要
        """
        self.api_key = api_key
        self.api_url = "https://api.deepseek.com/v1/chat/completions"
        
        # 验证参数合法性
        if not api_key:
            print("警告: 未提供API密钥，部分功能可能无法使用")
    
    def generate_text(self, prompt, max_tokens=500, temperature=0.7):
        """
        生成文本响应
        
        参数:
            prompt (str): 输入提示文本
            max_tokens (int): 生成的最大token数量
            temperature (float): 生成的随机性程度，值越高越随机
            
        返回:
            str: 生成的文本响应
        """
        if not prompt:
            return "错误: 提示文本不能为空"
            
        return self._generate_api(prompt, max_tokens, temperature)
    
    def _generate_api(self, prompt, max_tokens, temperature):
        """使用API生成文本"""
        if not self.api_key:
            return "错误: 使用API需要提供API密钥"
        
        headers = {
            "Content-Type": "application/json",
            "Authorization": f"Bearer {self.api_key}"
        }
        
        data = {
            "model": "deepseek-r1",
            "messages": [{"role": "user", "content": prompt}],
            "max_tokens": max_tokens,
            "temperature": temperature
        }
        
        try:
            response = requests.post(self.api_url, headers=headers, json=data, timeout=30)
            
            if response.status_code == 200:
                return response.json()["choices"][0]["message"]["content"]
            elif response.status_code == 401:
                return "API密钥无效或已过期，请检查API密钥"
            elif response.status_code == 429:
                return "API请求次数超限，请稍后再试"
            else:
                return f"API请求错误: {response.status_code}, {response.text}"
        except requests.exceptions.Timeout:
            return "API请求超时，请检查网络连接"
        except requests.exceptions.ConnectionError:
            return "无法连接到API服务器，请检查网络连接"
        except Exception as e:
            return f"API请求异常: {str(e)}"
    
    def chat(self, messages, max_tokens=500, temperature=0.7):
        """
        进行多轮对话
        
        参数:
            messages (list): 消息列表，格式为[{"role": "user", "content": "你好"}, ...]
            max_tokens (int): 生成的最大token数量
            temperature (float): 生成的随机性程度，值越高越随机
            
        返回:
            str: 生成的文本响应
        """
        if not messages or not isinstance(messages, list):
            return "错误: 消息必须是一个有效的列表"
            
        if not self.api_key:
            return "错误: 使用API需要提供API密钥"
        
        headers = {
            "Content-Type": "application/json",
            "Authorization": f"Bearer {self.api_key}"
        }
        
        data = {
            "model": "deepseek-r1",
            "messages": messages,
            "max_tokens": max_tokens,
            "temperature": temperature
        }
        
        try:
            response = requests.post(self.api_url, headers=headers, json=data, timeout=30)
            
            if response.status_code == 200:
                return response.json()["choices"][0]["message"]["content"]
            elif response.status_code == 401:
                return "API密钥无效或已过期，请检查API密钥"
            elif response.status_code == 429:
                return "API请求次数超限，请稍后再试"
            else:
                return f"API请求错误: {response.status_code}, {response.text}"
        except requests.exceptions.Timeout:
            return "API请求超时，请检查网络连接"
        except requests.exceptions.ConnectionError:
            return "无法连接到API服务器，请检查网络连接"
        except Exception as e:
            return f"API请求异常: {str(e)}"
    
    def summarize(self, text, max_tokens=200):
        """
        文本摘要功能
        
        参数:
            text (str): 要总结的文本
            max_tokens (int): 摘要的最大token数量
            
        返回:
            str: 文本摘要
        """
        if not text:
            return "错误: 要摘要的文本不能为空"
            
        prompt = f"请对以下文本进行总结，保持简洁明了：\n\n{text}"
        return self.generate_text(prompt, max_tokens, temperature=0.3)
    
    def translate(self, text, target_lang="中文"):
        """
        文本翻译功能
        
        参数:
            text (str): 要翻译的文本
            target_lang (str): 目标语言，默认为中文
            
        返回:
            str: 翻译后的文本
        """
        if not text:
            return "错误: 要翻译的文本不能为空"
            
        prompt = f"请将以下文本翻译成{target_lang}：\n\n{text}"
        return self.generate_text(prompt, max_tokens=len(text)*2, temperature=0.3)
    
    def answer_question(self, question, context=None):
        """
        回答问题功能
        
        参数:
            question (str): 问题
            context (str, 可选): 提供给模型的上下文信息
            
        返回:
            str: 回答
        """
        if not question:
            return "错误: 问题不能为空"
            
        if context:
            prompt = f"基于以下信息回答问题:\n\n背景信息: {context}\n\n问题: {question}"
        else:
            prompt = question
        
        return self.generate_text(prompt, max_tokens=500, temperature=0.7)


# 简单使用示例
if __name__ == "__main__":
    # 使用API方式（需要API密钥）
    # deepseek = DeepSeekR1(api_key="your-api-key-here")
    
    # 临时测试示例
    print("这是DeepSeek-R1助手类，请在实际项目中使用")
